#include "rclcpp/rclcpp.hpp"
#include <serial_driver/serial_driver.hpp>
#include <vector>
#include <cstdint>
#include <thread>
#include <atomic>
#include <mutex>
#include <memory>
#include <asio/serial_port.hpp>
#include "geometry_msgs/msg/vector3_stamped.hpp"
#include "geometry_msgs/msg/twist_stamped.hpp"
#include "sensor_msgs/msg/imu.hpp"
#include "nav_msgs/msg/odometry.hpp"
#include "tf2/LinearMath/Quaternion.h"
#include "tf2_ros/transform_broadcaster.h"
#include "tf2_ros/static_transform_broadcaster.h"
#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>
#include <Eigen/Dense>
#include <cmath>

using namespace drivers::serial_driver;

// 协议常量
constexpr uint8_t FRAME_HEADER = 0x55;
constexpr uint8_t FRAME_TAIL = 0xBB;
constexpr int MAX_FRAME_SIZE = 260;

// 命令码
enum Command : uint8_t {
  CMD_IMU_DATA = 0xA0,
  CMD_SPEED_DATA = 0xA1,
  CMD_BATTERY = 0xA3
};

class Stm32Bridge : public rclcpp::Node {
public:
  Stm32Bridge() : Node("stm32_bridge_node"), 
                  running_(true),
                  frame_state_(HEADER),
                  orientation_(Eigen::Quaterniond::Identity()),
                  gyro_bias_(Eigen::Vector3d::Zero()),
                  last_odom_update_(0, 0) {
    // 参数声明
    declare_parameter<std::string>("device_name", "/dev/ttyUSB0");
    declare_parameter<int>("baud_rate", 115200);
    declare_parameter<std::string>("odom_frame", "odom");
    declare_parameter<std::string>("base_frame", "base_link");
    declare_parameter<std::string>("imu_frame", "imu_link"); // 新增IMU frame参数
    declare_parameter<double>("odom_update_rate", 50.0);  // Hz
    declare_parameter<double>("gyro_bias_x", 0.0);  // 陀螺仪偏置
    declare_parameter<double>("gyro_bias_y", 0.0);
    declare_parameter<double>("gyro_bias_z", 0.0);

    // 获取参数
    odom_frame_ = get_parameter("odom_frame").as_string();
    base_frame_ = get_parameter("base_frame").as_string();
    imu_frame_ = get_parameter("imu_frame").as_string(); // 获取IMU frame
    double odom_update_rate = get_parameter("odom_update_rate").as_double();
    
    // 获取陀螺仪偏置参数
    gyro_bias_ = Eigen::Vector3d(
      get_parameter("gyro_bias_x").as_double(),
      get_parameter("gyro_bias_y").as_double(),
      get_parameter("gyro_bias_z").as_double()
    );

    // 初始化发布器
    imu_pub_ = this->create_publisher<sensor_msgs::msg::Imu>("imu", 10);
    odom_pub_ = this->create_publisher<nav_msgs::msg::Odometry>("odom", 10);
    tf_broadcaster_ = std::make_unique<tf2_ros::TransformBroadcaster>(*this);
    static_tf_broadcaster_ = std::make_shared<tf2_ros::StaticTransformBroadcaster>(*this); // 初始化静态TF广播器
    
    RCLCPP_INFO(get_logger(), "节点已初始化，里程计更新频率: %.1f Hz", odom_update_rate);
    RCLCPP_INFO(get_logger(), "陀螺仪偏置: X=%.4f, Y=%.4f, Z=%.4f", 
                gyro_bias_.x(), gyro_bias_.y(), gyro_bias_.z());

    // 初始化定时器更新里程计
    auto update_period = std::chrono::duration<double>(1.0 / odom_update_rate);
    odom_timer_ = this->create_wall_timer(
      std::chrono::duration_cast<std::chrono::nanoseconds>(update_period),
      std::bind(&Stm32Bridge::odom_update_callback, this));

    // 初始化串口
    init_serial();
    
    // 发布一次静态TF（base_link到imu_link）
    publish_static_tf();
  }

  ~Stm32Bridge() {
    running_ = false;
    if (receiver_thread_.joinable()) {
      receiver_thread_.join();
    }

    if (serial_port_ && serial_port_->is_open()) {
      serial_port_->close();
    }
    RCLCPP_INFO(get_logger(), "串口资源已释放");
  }

private:
  void init_serial() {
    try {
      // 获取参数
      const std::string device_name = get_parameter("device_name").as_string();
      const uint32_t baud_rate = static_cast<uint32_t>(get_parameter("baud_rate").as_int());

      // 创建IO上下文
      io_context_ = std::make_shared<IoContext>();

      // 配置串口参数
      const FlowControl fc = FlowControl::NONE;
      const Parity parity = Parity::NONE;
      const StopBits stop_bits = StopBits::ONE;

      serial_port_ = std::make_unique<SerialPort>(
        *io_context_,
        device_name,
        SerialPortConfig{baud_rate, fc, parity, stop_bits}
      );

      // 打开端口
      serial_port_->open();

      RCLCPP_INFO(get_logger(), "串口初始化成功: %s @ %u bps",
                 device_name.c_str(), baud_rate);

      // 启动接收线程
      receiver_thread_ = std::thread(&Stm32Bridge::receive_loop, this);
    } catch (const std::exception &ex) {
      RCLCPP_FATAL(get_logger(), "串口初始化失败: %s", ex.what());
      rclcpp::shutdown();
    }
  }

  void receive_loop() {
    while (running_ && rclcpp::ok()) {
      try {
        // 一次读取一个字节
        std::vector<uint8_t> buffer(1);
        const size_t bytes_read = serial_port_->receive(buffer);

        if (bytes_read == 1) {
          std::lock_guard<std::mutex> lock(frame_mutex_);
          process_byte(buffer[0]);
        }
      } catch (const std::exception& e) {
        RCLCPP_ERROR(get_logger(), "接收错误: %s", e.what());
      }
    }
  }

  void process_byte(uint8_t byte) {
    // 状态机重置条件：检测到连续帧头
    if (frame_buffer_.size() == 0 && byte == FRAME_HEADER) {
      frame_state_ = HEADER;
    } else if (frame_state_ == HEADER && byte == FRAME_HEADER) {
      frame_buffer_.clear();
      frame_buffer_.push_back(FRAME_HEADER);
      frame_buffer_.push_back(FRAME_HEADER);
      calc_checksum_ = 0;
      frame_state_ = COMMAND;
      return;
    }

    // 状态机处理
    switch (frame_state_) {
      case HEADER:
        if (byte == FRAME_HEADER) {
          frame_buffer_.push_back(byte);
          frame_state_ = HEADER2; // 新增 HEADER2 状态
        } else {
          reset_frame_state();
        }
        break;

      case HEADER2:
        if (byte == FRAME_HEADER) {
          frame_buffer_.push_back(byte);
          frame_state_ = COMMAND;
          calc_checksum_ = 0;
        } else {
          reset_frame_state();
        }
        break;

      case COMMAND:
        current_command_ = byte;
        frame_buffer_.push_back(byte);
        calc_checksum_ += byte;
        frame_state_ = LENGTH;
        break;

      case LENGTH:
        data_length_ = byte;
        frame_buffer_.push_back(byte);
        calc_checksum_ += byte;
        data_counter_ = 0;
        frame_state_ = (data_length_ > 0) ? DATA : CHECKSUM;
        break;

      case DATA:
        frame_buffer_.push_back(byte);
        calc_checksum_ += byte;
        if (++data_counter_ >= data_length_) {
          frame_state_ = CHECKSUM;
        }
        break;

      case CHECKSUM:
        if (byte == calc_checksum_) {
          frame_buffer_.push_back(byte);
          frame_state_ = TAIL;
        } else {
          RCLCPP_WARN(get_logger(), "校验错误: 预期0x%02X 实际0x%02X",
                     calc_checksum_, byte);
          reset_frame_state();
        }
        break;

      case TAIL:
        if (byte == FRAME_TAIL) {
          frame_buffer_.push_back(byte);
          parse_frame();  // 处理完整帧
        } else {
          RCLCPP_WARN(get_logger(), "帧尾错误: 预期0x%02X 实际0x%02X",
                     FRAME_TAIL, byte);
        }
        reset_frame_state();
        break;
    }
  }

  void reset_frame_state() {
    frame_buffer_.clear();
    frame_state_ = HEADER;
    calc_checksum_ = 0;
    data_length_ = 0;
    data_counter_ = 0;
  }

  void parse_frame() {
    // 最小帧长度检查
    if (frame_buffer_.size() < static_cast<size_t>(6 + data_length_)) {
      RCLCPP_WARN(get_logger(), "帧长度不足");
      return;
    }

    const uint8_t* data = frame_buffer_.data() + 4;  // 跳过头部（2个0x55）+ command + length
    auto current_time = this->now();

    switch (current_command_) {
      case CMD_IMU_DATA:
        if (data_length_ == 28) {
          float ax = read_float_le(data + 4);
          float ay = read_float_le(data + 8);
          float az = read_float_le(data + 12);
          float gx = read_float_le(data + 16);
          float gy = read_float_le(data + 20);
          float gz = read_float_le(data + 24);

          {
            std::lock_guard<std::mutex> lock(imu_mutex_);
            // 存储IMU原始数据
            last_accel_ = Eigen::Vector3d(ax, ay, az);
            last_gyro_ = Eigen::Vector3d(gx, gy, gz);
            last_imu_time_ = current_time;
          }

          // 处理IMU数据并发布
          process_imu_data(current_time);

          RCLCPP_DEBUG(get_logger(), "IMU数据: Accel(%.2f,%.2f,%.2f) Gyro(%.2f,%.2f,%.2f)",
                       ax, ay, az, gx, gy, gz);
        } else {
          RCLCPP_WARN(get_logger(), "IMU数据长度错误: %d != 28", data_length_);
        }
        break;

      case CMD_SPEED_DATA:
        if (data_length_ == 6) {
          int16_t vx = read_i16_be(data);
          int16_t vy = read_i16_be(data + 2);
          int16_t vz = read_i16_be(data + 4);

          {
            std::lock_guard<std::mutex> lock(odom_mutex_);
            // 更新速度数据
            current_velocity_.linear.x = vx / 1000.0;
            current_velocity_.linear.y = vy / 1000.0;
            current_velocity_.angular.z = vz / 1000.0;
            last_vel_time_ = current_time;
          }

          RCLCPP_DEBUG(get_logger(), "速度: Vx=%.3f m/s, Vy=%.3f m/s, Vz=%.3f rad/s",
                       vx / 1000.0, vy / 1000.0, vz / 1000.0);
        } else {
          RCLCPP_WARN(get_logger(), "速度数据长度错误: %d != 6", data_length_);
        }
        break;

      case CMD_BATTERY:
        if (data_length_ == 2) {
          uint16_t voltage = read_u16_le(data);
          RCLCPP_INFO_THROTTLE(
            get_logger(), *get_clock(), 5000,
            "电池电压: %.2f V", voltage / 1000.0);
        } else {
          RCLCPP_WARN(get_logger(), "电池数据长度错误: %d != 2", data_length_);
        }
        break;

      default:
        RCLCPP_WARN(get_logger(), "未知命令码: 0x%02X", current_command_);
        break;
    }
  }

  void process_imu_data(rclcpp::Time timestamp) {
    Eigen::Vector3d accel, gyro;
    
    {
      std::lock_guard<std::mutex> lock(imu_mutex_);
      accel = last_accel_;
      gyro = last_gyro_;
    }

    // 应用低通滤波器
    static Eigen::Vector3d last_filtered_accel = accel;
    static Eigen::Vector3d last_filtered_gyro = gyro;
    static const double alpha = 0.2;
    
    Eigen::Vector3d filtered_accel = alpha * accel + (1 - alpha) * last_filtered_accel;
    Eigen::Vector3d filtered_gyro = alpha * gyro + (1 - alpha) * last_filtered_gyro;
    
    last_filtered_accel = filtered_accel;
    last_filtered_gyro = filtered_gyro;

    // 移除陀螺仪偏置
    Eigen::Vector3d gyro_corrected = filtered_gyro - gyro_bias_;

    // 姿态更新
    if (filtered_accel.norm() > 0.1) {  // 防止在失重状态下计算
      Eigen::Vector3d down = filtered_accel.normalized();
      
      // 使用重力向量计算姿态
      Eigen::Quaterniond accel_quat;
      accel_quat.setFromTwoVectors(Eigen::Vector3d::UnitZ(), down);
      
      // 使用陀螺仪积分
      double dt = 0.01; // 假设固定时间间隔
      Eigen::Quaterniond gyro_quat(1, 
                                  gyro_corrected.x() * dt / 2, 
                                  gyro_corrected.y() * dt / 2, 
                                  gyro_corrected.z() * dt / 2);
      gyro_quat.normalize();
      
      // 互补滤波融合
      double alpha = 0.98; // 陀螺仪权重
      orientation_ = orientation_.slerp(alpha, accel_quat * orientation_ * gyro_quat);
      orientation_.normalize();
    }

    // 创建IMU消息
    auto imu_msg = sensor_msgs::msg::Imu();
    imu_msg.header.stamp = timestamp;
    imu_msg.header.frame_id = imu_frame_; // 使用参数设置的frame_id
    
    // 设置协方差矩阵（示例值，需根据传感器标定调整）
    imu_msg.orientation_covariance = {
        0.05, 0.0,  0.0,   // 姿态协方差 (3x3 矩阵)
        0.0,  0.05, 0.0,
        0.0,  0.0,  0.05
    };
    imu_msg.angular_velocity_covariance = {
        0.002, 0.0,   0.0,   // 角速度协方差
        0.0,   0.002, 0.0,
        0.0,   0.0,   0.002
    };
    imu_msg.linear_acceleration_covariance = {
        0.03, 0.0,  0.0,   // 加速度协方差
        0.0,  0.03, 0.0,
        0.0,  0.0,  0.03
    };

    // 设置方向
    imu_msg.orientation.x = orientation_.x();
    imu_msg.orientation.y = orientation_.y();
    imu_msg.orientation.z = orientation_.z();
    imu_msg.orientation.w = orientation_.w();
    
    // 设置角速度
    imu_msg.angular_velocity.x = gyro_corrected.x();
    imu_msg.angular_velocity.y = gyro_corrected.y();
    imu_msg.angular_velocity.z = gyro_corrected.z();
    
    // 设置线性加速度
    imu_msg.linear_acceleration.x = filtered_accel.x();
    imu_msg.linear_acceleration.y = filtered_accel.y();
    imu_msg.linear_acceleration.z = filtered_accel.z();
    
    imu_pub_->publish(imu_msg);
  }

  void publish_static_tf() {
    // 发布base_link到imu_link的静态变换
    geometry_msgs::msg::TransformStamped imu_tf;
    imu_tf.header.stamp = this->now();
    imu_tf.header.frame_id = base_frame_;    // base_link
    imu_tf.child_frame_id = imu_frame_;      // imu_link
    
    // 设置IMU相对于base_link的变换（根据实际安装位置调整）
    imu_tf.transform.translation.x = 0.0;
    imu_tf.transform.translation.y = 0.0;
    imu_tf.transform.translation.z = 0.1; // 假设IMU安装在base_link上方10cm处
    
    // 设置方向（根据实际安装方向调整）
    tf2::Quaternion q_imu;
    q_imu.setRPY(0, 0, 0); // 假设IMU与base_link方向一致
    imu_tf.transform.rotation = tf2::toMsg(q_imu);
    
    // 发布静态TF（只发布一次）
    static_tf_broadcaster_->sendTransform(imu_tf);
    RCLCPP_INFO(get_logger(), "已发布静态TF: %s -> %s", 
                base_frame_.c_str(), imu_frame_.c_str());
  }

  void odom_update_callback() {
    geometry_msgs::msg::Twist current_vel;
    rclcpp::Time vel_time;
    
    {
      std::lock_guard<std::mutex> lock(odom_mutex_);
      current_vel = current_velocity_;
      vel_time = last_vel_time_;
    }

    rclcpp::Time current_time = this->now();
    
    // 第一次调用初始化
    if (last_odom_update_.nanoseconds() == 0) {
      last_odom_update_ = current_time;
      return;
    }
    
    // 计算时间差
    rclcpp::Duration dt_duration = current_time - last_odom_update_;
    double dt = dt_duration.seconds();
    last_odom_update_ = current_time;
    
    // 更新位置（欧拉积分）
    double vx = current_vel.linear.x;
    double vy = current_vel.linear.y;
    double wz = current_vel.angular.z;
    
    // 计算位置增量
    double delta_x = vx * dt;
    double delta_y = vy * dt;
    
    // 更新位置
    x_ += delta_x;
    y_ += delta_y;
    theta_ += wz * dt; // 更新偏航角
    
    // 创建四元数
    tf2::Quaternion q;
    q.setRPY(0, 0, theta_);
    
    // 创建里程计消息
    auto odom_msg = nav_msgs::msg::Odometry();
    odom_msg.header.stamp = current_time;
    odom_msg.header.frame_id = odom_frame_;
    odom_msg.child_frame_id = base_frame_;
    
    // 设置位姿协方差 (6x6 矩阵，行优先)
    // 注意: 协方差矩阵是36个元素的一维数组，按行优先存储
    odom_msg.pose.covariance[0]  = 0.1;   // x方差
    odom_msg.pose.covariance[7]  = 0.1;   // y方差
    odom_msg.pose.covariance[14] = 0.01;  // z方差
    odom_msg.pose.covariance[21] = 0.05;  // roll方差
    odom_msg.pose.covariance[28] = 0.05;  // pitch方差
    odom_msg.pose.covariance[35] = 0.01;  // yaw方差
    
    // 设置速度协方差
    odom_msg.twist.covariance[0]  = 0.02;  // vx方差
    odom_msg.twist.covariance[7]  = 0.02;  // vy方差
    odom_msg.twist.covariance[14] = 0.01;  // vz方差
    odom_msg.twist.covariance[21] = 0.01;  // wx方差
    odom_msg.twist.covariance[28] = 0.01;  // wy方差
    odom_msg.twist.covariance[35] = 0.03;  // wz方差

    // 位置
    odom_msg.pose.pose.position.x = x_;
    odom_msg.pose.pose.position.y = y_;
    odom_msg.pose.pose.position.z = 0.0;
    odom_msg.pose.pose.orientation.x = q.x();
    odom_msg.pose.pose.orientation.y = q.y();
    odom_msg.pose.pose.orientation.z = q.z();
    odom_msg.pose.pose.orientation.w = q.w();
    
    // 速度
    odom_msg.twist.twist = current_vel;
    
    odom_pub_->publish(odom_msg);
    
    // 发布TF变换：odom->base_link
    geometry_msgs::msg::TransformStamped odom_tf;
    odom_tf.header.stamp = current_time;
    odom_tf.header.frame_id = odom_frame_;
    odom_tf.child_frame_id = base_frame_;
    odom_tf.transform.translation.x = x_;
    odom_tf.transform.translation.y = y_;
    odom_tf.transform.translation.z = 0.0;
    odom_tf.transform.rotation = tf2::toMsg(q);
    
    tf_broadcaster_->sendTransform(odom_tf);
    
    RCLCPP_DEBUG(get_logger(), "里程计: X=%.2f, Y=%.2f, Theta=%.2f", x_, y_, theta_);
  }

  // 小端读取辅助函数
  uint16_t read_u16_le(const uint8_t* data) {
    return static_cast<uint16_t>((data[1] << 8) | data[0]);
  }

  int16_t read_i16_le(const uint8_t* data) {
    return static_cast<int16_t>((data[1] << 8) | data[0]);
  }

  uint32_t read_u32_le(const uint8_t* data) {
    return (static_cast<uint32_t>(data[3]) << 24) |
           (static_cast<uint32_t>(data[2]) << 16) |
           (static_cast<uint32_t>(data[1]) << 8) |
           data[0];
  }

  int16_t read_i16_be(const uint8_t* data) {
    return static_cast<int16_t>((data[0] << 8) | data[1]); // 高字节在前
  }

  float read_float_le(const uint8_t* data) {
    uint32_t tmp = (static_cast<uint32_t>(data[3]) << 24) |
                   (static_cast<uint32_t>(data[2]) << 16) |
                   (static_cast<uint32_t>(data[1]) << 8) |
                   data[0];
    return *reinterpret_cast<float*>(&tmp);
  }

    // 线程同步
  std::mutex frame_mutex_;
  std::mutex imu_mutex_;
  std::mutex odom_mutex_;
  std::atomic<bool> running_;

  // 帧解析状态
  enum FrameState {
    HEADER,
    HEADER2,
    COMMAND,
    LENGTH,
    DATA,
    CHECKSUM,
    TAIL
  } frame_state_;

  // 串口资源
  std::shared_ptr<IoContext> io_context_;
  std::unique_ptr<SerialPort> serial_port_;
  std::thread receiver_thread_;

  // ROS 发布器
  rclcpp::Publisher<sensor_msgs::msg::Imu>::SharedPtr imu_pub_;
  rclcpp::Publisher<nav_msgs::msg::Odometry>::SharedPtr odom_pub_;
  std::unique_ptr<tf2_ros::TransformBroadcaster> tf_broadcaster_;
  std::shared_ptr<tf2_ros::StaticTransformBroadcaster> static_tf_broadcaster_; // 静态TF广播器

  // 帧处理状态
  uint8_t current_command_;
  uint8_t data_length_;
  uint16_t data_counter_;
  uint8_t calc_checksum_;
  std::vector<uint8_t> frame_buffer_;

  // IMU处理相关
  Eigen::Quaterniond orientation_;
  Eigen::Vector3d gyro_bias_;
  Eigen::Vector3d last_accel_;
  Eigen::Vector3d last_gyro_;
  rclcpp::Time last_imu_time_{0, 0};

  // 里程计相关
  double x_ = 0.0;
  double y_ = 0.0;
  double theta_ = 0.0; // 偏航角
  geometry_msgs::msg::Twist current_velocity_;
  rclcpp::Time last_vel_time_{0, 0};
  rclcpp::Time last_odom_update_{0, 0};
  std::string odom_frame_;
  std::string base_frame_;
  std::string imu_frame_; // IMU坐标系名称

  // 定时器
  rclcpp::TimerBase::SharedPtr odom_timer_;
};

int main(int argc, char **argv) {
  rclcpp::init(argc, argv);
  auto node = std::make_shared<Stm32Bridge>();
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}